import os
import cv2
import numpy as np
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras.utils import to_categorical
import tensorflow as tf

class DataProcessor:
    def __init__(self, data_dir=None):
        """初始化数据处理器"""
        self.data_dir = data_dir
        self.target_size = (224, 224)
        self.img_size = 224
    
    def preprocess_single_image(self, image_path):
        """预处理单张图片"""
        try:
            # 读取图片
            img = tf.keras.preprocessing.image.load_img(
                image_path,
                target_size=(self.img_size, self.img_size)
            )
            
            # 转换为数组
            img_array = tf.keras.preprocessing.image.img_to_array(img)
            
            # 归一化
            img_array = img_array / 255.0
            
            return img_array
            
        except Exception as e:
            raise Exception(f"图片预处理失败: {str(e)}")
    
    def load_and_preprocess(self):
        """加载和预处理数据集"""
        if self.data_dir is None:
            raise ValueError("未指定数据集目录")
        
        X = []
        y = []
        class_names = []
        
        # 处理目标物品图片
        target_dir = os.path.join(self.data_dir, 'target')
        if os.path.exists(target_dir):
            class_names.append('target')
            for img_name in os.listdir(target_dir):
                if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                    try:
                        img_path = os.path.join(target_dir, img_name)
                        img_array = self.preprocess_single_image(img_path)
                        X.append(img_array)
                        y.append(1)  # 目标物品标签为1
                    except Exception as e:
                        print(f"处理图片失败 {img_path}: {str(e)}")
        
        # 处理背景图片
        background_dir = os.path.join(self.data_dir, 'background')
        if os.path.exists(background_dir):
            class_names.append('background')
            for img_name in os.listdir(background_dir):
                if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                    try:
                        img_path = os.path.join(background_dir, img_name)
                        img_array = self.preprocess_single_image(img_path)
                        X.append(img_array)
                        y.append(0)  # 背景图片标签为0
                    except Exception as e:
                        print(f"处理图片失败 {img_path}: {str(e)}")
        
        if not X:
            raise ValueError("没有找到有效的图片")
        
        # 转换为numpy数组
        X = np.array(X)
        y = np.array(y)
        
        return X, y, class_names
    
    def _process_image(self, img_path):
        """处理单张图片"""
        try:
            # 读取图片
            img = cv2.imread(img_path)
            if img is None:
                return None
            
            # BGR转RGB
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 调整大小
            img = cv2.resize(img, (self.img_size, self.img_size))
            
            # 归一化 - 使用ResNet的预处理方法
            img = img.astype(np.float32)
            img = tf.keras.applications.resnet_v2.preprocess_input(img)
            
            return img
        except Exception as e:
            print(f"处理图片失败 {img_path}: {str(e)}")
            return None
    
    def _encode_labels(self, labels):
        """编码标签"""
        if not labels:
            raise ValueError("没有标签数据可供编码")
        
        # 将文本标签转换为数字
        numeric_labels = self.label_encoder.fit_transform(labels)
        
        # 根据类别数量决定编码方式
        num_classes = len(self.label_encoder.classes_)
        print(f"\n检测到 {num_classes} 个类别:")
        for i, class_name in enumerate(self.label_encoder.classes_):
            print(f"- 类别 {i}: {class_name}")
        
        if num_classes == 2:
            # 二分类：直接返回 0/1 标签
            return numeric_labels
        else:
            # 多分类：转换为 one-hot 编码
            return to_categorical(numeric_labels)